{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/peremartra/Large-Language-Model-Notebooks-Course/blob/main/6-PRUNING/6_6_pruning_attention_layers.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LV5KPbfseL8S"
      },
      "source": [
        "<div>\n",
        "    <h1>Large Language Models Projects</a></h1>\n",
        "    <h3>Apply and Implement Strategies for Large Language Models</h3>\n",
        "    <h2>Pruning Attention Layers</h2>\n",
        "    <h3>Not All Attention is needed</h3>\n",
        "</div>\n",
        "\n",
        "by [Pere Martra](https://www.linkedin.com/in/pere-martra/)\n",
        "\n",
        "_______\n",
        "Models: meta-llama/Llama-3.2\n",
        "\n",
        "Colab Environment: GPU L4 for 3B Models\n",
        "\n",
        "T4 for 1B Model.\n",
        "\n",
        "Keys:\n",
        "* Pruning\n",
        "* Attention\n",
        "\n",
        "References:\n",
        "* [What Matters in Transformers? Not All Attention is Needed](https://arxiv.org/abs/2406.15786)\n",
        "* [Resource-Efficient Transformer Pruning for Finetuning of Large Models](https://openaccess.thecvf.com/content/CVPR2024/html/Ilhan_Resource-Efficient_Transformer_Pruning_for_Finetuning_of_Large_Models_CVPR_2024_paper.html)\n",
        "\n",
        "_______\n",
        "**disclaimer: The pruning / knowledge distillation section has been created after the first edition of the book was published. They are not included in the book’s original content but are intended to supplement and expand on the topics covered.**\n",
        "\n",
        "This is the unofficial repository for the book:\n",
        "        <a href=\"https://amzn.to/4eanT1g\"> <b>Large Language Models:</b> Apply and Implement Strategies for Large Language Models</a> (Apress).\n",
        "        The book is based on the content of this repository, but the notebooks are being updated, and I am incorporating new examples and chapters.\n",
        "        If you are looking for the official repository for the book, with the original notebooks, you should visit the\n",
        "        <a href=\"https://github.com/Apress/Large-Language-Models-Projects\">Apress repository</a>, where you can find all the notebooks in their original format as they appear in the book.\n",
        "\n",
        "______\n",
        "# Introduction\n",
        "This notebook implements the paper: [What Matters in Transformers? Not all Attention is Needed](https://arxiv.org/abs/2406.15786).\n",
        "Although I followed the paper's guidelines, I made some adjustments to make the code clearer and easier to understand.\n",
        "\n",
        "The original paper demonstrates that larger models tend to have excessive redundancy in their attention layers. They achieved a 48.4% increase in inference performance for a Llama-2-70B model with only a minor 2.4% drop in response quality, **just bypassing the 50% of the Attention layers!**\n",
        "\n",
        "In this notebook, tests have been conducted using Llama-3.2-1B and 3B models. With these models, I found that removing 50% of the attention layers significantly impacted the model's functionality. However, the 3B model handled the removal of these layers much better. This suggests that redundancy may become more pronounced as model size increases.\n",
        "\n",
        "# Methodology.\n",
        "To identify which layers contribute the least, the cosine distance between the layer's input and output is measured. In the paper, this distance is calculated using a test dataset, while in the notebook, I used a simple prompt to activate the layers.  \n",
        "\n",
        "This method of measuring the importance of attention layers and their contribution to the model allows pruning to be tailored to a specific dataset. This approach can lead to the creation of more efficient models for specialized sectors such as healthcare or finance.\n",
        "\n",
        "\n",
        "Once the layer contributing the least to the final output is identified (the one with the smallest difference between input and output), it is added to a list included in the configuration file.\n",
        "\n",
        "This list is then referenced during inference by a new forward function that replaces the original one for attention layers. When this new function detects that a layer is in the list, it skips its execution and simply returns the input without modifications.\n",
        "\n",
        "The process of identifying layers to deactivate and marking them as non-executable is one-shot. In other words, it does  determine all the layers to skip in one go, as recommended in the paper.\n",
        "\n",
        "The iterative implementation has a significant drawback: the test dataset must be processed for each layer to be deactivated. The paper's authors note that while the iterative method may bring slight improvements, the added computational cost is not justified. However, since this is an example notebook, and there is no test dataset—just a small prompt—and the layer selection process takes only seconds, I chose the iterative approach.\n",
        "\n",
        "This pruning method does not produce a smaller model, as the layers are not physically removed. They remain in the model but are not executed, resulting in improved inference response times.\n",
        "______"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DwYeKwswnkTG"
      },
      "source": [
        "# Install libraries & Configure variables."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PblPrYCiYTl8",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "5374f0f7-2a56-4ce4-ea9d-d6e4fcaacfeb"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m363.4/363.4 MB\u001b[0m \u001b[31m2.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m13.8/13.8 MB\u001b[0m \u001b[31m82.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m24.6/24.6 MB\u001b[0m \u001b[31m94.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m883.7/883.7 kB\u001b[0m \u001b[31m49.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m664.8/664.8 MB\u001b[0m \u001b[31m1.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m211.5/211.5 MB\u001b[0m \u001b[31m10.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m56.3/56.3 MB\u001b[0m \u001b[31m38.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m127.9/127.9 MB\u001b[0m \u001b[31m17.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m207.5/207.5 MB\u001b[0m \u001b[31m2.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m21.1/21.1 MB\u001b[0m \u001b[31m103.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m491.5/491.5 kB\u001b[0m \u001b[31m36.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m193.6/193.6 kB\u001b[0m \u001b[31m18.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "gcsfs 2025.3.2 requires fsspec==2025.3.2, but you have fsspec 2025.3.0 which is incompatible.\u001b[0m\u001b[31m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m50.5/50.5 kB\u001b[0m \u001b[31m4.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m51.8/51.8 kB\u001b[0m \u001b[31m5.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.9/3.9 MB\u001b[0m \u001b[31m103.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m84.0/84.0 kB\u001b[0m \u001b[31m7.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m243.3/243.3 kB\u001b[0m \u001b[31m24.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m104.1/104.1 kB\u001b[0m \u001b[31m11.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m91.1/91.1 kB\u001b[0m \u001b[31m5.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Building wheel for rouge-score (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Building wheel for sqlitedict (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Building wheel for word2number (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "Collecting hf_xet\n",
            "  Downloading hf_xet-1.1.2-cp37-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (879 bytes)\n",
            "Downloading hf_xet-1.1.2-cp37-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (5.2 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m5.2/5.2 MB\u001b[0m \u001b[31m101.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hInstalling collected packages: hf_xet\n",
            "Successfully installed hf_xet-1.1.2\n"
          ]
        }
      ],
      "source": [
        "!pip install -q torch==2.6.0\n",
        "!pip install -q torchvision==0.21.0\n",
        "!pip install -q transformers==4.51.3\n",
        "!pip install -q datasets==3.6.0\n",
        "!pip install -q lm-eval==0.4.8\n",
        "\n",
        "!pip install hf_xet #To speed up downloads from HF."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6qN0mu6IHqpy"
      },
      "outputs": [],
      "source": [
        "import logging\n",
        "import math\n",
        "import os\n",
        "import sys\n",
        "import shutil\n",
        "from copy import deepcopy\n",
        "\n",
        "import torch\n",
        "import torch.nn.functional as F\n",
        "from transformers import AutoModelForCausalLM, AutoTokenizer\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J8iRr5iapy5q"
      },
      "source": [
        "# Download the Model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "u9IFTYa9P6Zy",
        "outputId": "06eea4ec-58bd-47dc-dab1-7a307d1cc09a"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Using device: cuda\n"
          ]
        }
      ],
      "source": [
        "# Check if GPU is available\n",
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "print(f\"Using device: {device}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sbnDW_tZRTGp"
      },
      "outputs": [],
      "source": [
        "#model_name = 'meta-llama/Llama-3.2-1B'\n",
        "model_name = 'meta-llama/Llama-3.2-3B'\n",
        "model = AutoModelForCausalLM.from_pretrained(model_name).to(device)\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
        "#tokenizer.pad_token = tokenizer.eos_token  # Set pad token"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w0ifeungrbHw"
      },
      "source": [
        "## Study the structure.\n",
        "* Llama-3.2-1B\n",
        "```\n",
        "LlamaForCausalLM(\n",
        "  (model): LlamaModel(\n",
        "    (embed_tokens): Embedding(128256, 2048)\n",
        "    (layers): ModuleList(\n",
        "      (0-15): 16 x LlamaDecoderLayer(\n",
        "        (self_attn): LlamaSdpaAttention(\n",
        "          (q_proj): Linear(in_features=2048, out_features=2048, bias=False)\n",
        "          (k_proj): Linear(in_features=2048, out_features=512, bias=False)\n",
        "          (v_proj): Linear(in_features=2048, out_features=512, bias=False)\n",
        "          (o_proj): Linear(in_features=2048, out_features=2048, bias=False)\n",
        "          (rotary_emb): LlamaRotaryEmbedding()\n",
        "        )\n",
        "        (mlp): LlamaMLP(\n",
        "          (gate_proj): Linear(in_features=2048, out_features=8192, bias=False)\n",
        "          (up_proj): Linear(in_features=2048, out_features=8192, bias=False)\n",
        "          (down_proj): Linear(in_features=8192, out_features=2048, bias=False)\n",
        "          (act_fn): SiLU()\n",
        "        )\n",
        "        (input_layernorm): LlamaRMSNorm((2048,), eps=1e-05)\n",
        "        (post_attention_layernorm): LlamaRMSNorm((2048,), eps=1e-05)\n",
        "      )\n",
        "    )\n",
        "    (norm): LlamaRMSNorm((2048,), eps=1e-05)\n",
        "    (rotary_emb): LlamaRotaryEmbedding()\n",
        "  )\n",
        "  (lm_head): Linear(in_features=2048, out_features=128256, bias=False)\n",
        ")\n",
        "```\n",
        "\n",
        "\n",
        "The model follows the typical structure of modern Llama models, consisting of blocks made up of an Attention layer and an MLP layer with a GLU structure.\n",
        "\n",
        "> If you want to see an example of how to perform pruning on the MLP layers of the model, you can check out the notebook:[Pruning Llama 3.2.](https://github.com/peremartra/Large-Language-Model-Notebooks-Course/blob/main/6-PRUNING/6_3_pruning_structured_llama3.2-1b_OK.ipynb) y leer el paper [Exploring GLU expansion ratios: Structured pruning in Llama-3.2 models](https://osf.io/preprints/osf/qgxea)\n",
        "\n",
        "\n",
        "\n",
        "Since the layers form a block, the attention layer cannot be removed without also removing the accompanying MLP layer. For this reason, the decision was made to deactivate their execution during inference.\n",
        "\n",
        "The 1B model has 16 layers, as shown in the structure above, while the 3B model has 28 layers.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Inference function & Test Base Model\n",
        "\n",
        "The `get_output` function is designed to generate text  and measure the time taken for different stages of the generation process.\n",
        "\n",
        "It provides insights into the performance of the model and can be used to evaluate the efficiency of text generation."
      ],
      "metadata": {
        "id": "vF4cHUb_rICs"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2igvy4z6rGgy"
      },
      "outputs": [],
      "source": [
        "import time\n",
        "\n",
        "def get_output(prompt, model=model, tokenizer=tokenizer, num_runs=1, max_length=50):\n",
        "    total_time = 0\n",
        "    generated_outputs = []\n",
        "\n",
        "    for run in range(num_runs):\n",
        "        # Start timing\n",
        "        start_time = time.time()\n",
        "\n",
        "        # Tokenization time\n",
        "        token_start = time.time()\n",
        "        inputs = tokenizer(prompt, return_tensors='pt').to(device)\n",
        "        token_time = time.time() - token_start\n",
        "\n",
        "        # Generation time\n",
        "        gen_start = time.time()\n",
        "        outputs = model.generate(\n",
        "            inputs['input_ids'],\n",
        "            attention_mask=inputs['attention_mask'],\n",
        "            max_length=max_length,\n",
        "            num_return_sequences=1,\n",
        "            pad_token_id=tokenizer.pad_token_id,\n",
        "            temperature=None,\n",
        "            top_p=None,\n",
        "            do_sample=False,  # Disable sampling\n",
        "            num_beams=5,      # Use beam search\n",
        "            early_stopping=True,  # Stop when end-of-sequence token is generated\n",
        "            no_repeat_ngram_size=2  # Prevent repetition of 2-grams\n",
        "        )\n",
        "        gen_time = time.time() - gen_start\n",
        "\n",
        "        # Decoding time\n",
        "        decode_start = time.time()\n",
        "        generated = tokenizer.decode(outputs[0], skip_special_tokens=True)\n",
        "        decode_time = time.time() - decode_start\n",
        "\n",
        "        # Total time for this run\n",
        "        total_time += time.time() - start_time\n",
        "        generated_outputs.append(generated)\n",
        "\n",
        "        if num_runs > 1:\n",
        "            print(f\"\\nRun {run + 1}:\")\n",
        "        print(f\"Tokenization time: {token_time*1000:.2f} ms\")\n",
        "        print(f\"Generation time: {gen_time*1000:.2f} ms\")\n",
        "        print(f\"Decoding time: {decode_time*1000:.2f} ms\")\n",
        "        print(f\"Total time: {(time.time() - start_time)*1000:.2f} ms\")\n",
        "\n",
        "    if num_runs > 1:\n",
        "        avg_time = total_time / num_runs\n",
        "        print(f\"\\nAverage time over {num_runs} runs: {avg_time*1000:.2f} ms\")\n",
        "\n",
        "    return generated_outputs[0] if num_runs == 1 else generated_outputs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "a836bc52-d0f6-418c-9c09-06866217807a",
        "id": "lH7cotAxrhO3"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Setting `pad_token_id` to `eos_token_id`:128001 for open-end generation.\n",
            "Setting `pad_token_id` to `eos_token_id`:128001 for open-end generation.\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Run 1:\n",
            "Tokenization time: 1.84 ms\n",
            "Generation time: 2701.53 ms\n",
            "Decoding time: 0.27 ms\n",
            "Total time: 2703.73 ms\n",
            "\n",
            "Run 2:\n",
            "Tokenization time: 0.55 ms\n",
            "Generation time: 1657.01 ms\n",
            "Decoding time: 0.21 ms\n",
            "Total time: 1657.85 ms\n",
            "\n",
            "Average time over 2 runs: 2180.71 ms\n",
            "Generated text: ['Paris is the capital of France. It is located in the north-central part of the country, on the river Seine. The city has a population of over 2 million people, making it the largest city in France and the second-largest city', 'Paris is the capital of France. It is located in the north-central part of the country, on the river Seine. The city has a population of over 2 million people, making it the largest city in France and the second-largest city']\n"
          ]
        }
      ],
      "source": [
        "# Test the original model\n",
        "prompt = \"Paris is the capital of\"\n",
        "generated = get_output(prompt, num_runs=2)\n",
        "print(f\"Generated text: {generated}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The text generation of the original model, as expected, works perfectly and returns a correct and meaningful sentence."
      ],
      "metadata": {
        "id": "mo4IjOYGry0W"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "model.to(\"cpu\")               # actual data moves ↙\n",
        "torch.cuda.empty_cache()      # allocator drops cached blocks"
      ],
      "metadata": {
        "id": "bLN1_gLdt7Rx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vjgf4WA_vF3B"
      },
      "source": [
        "# Pruning the Model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TQDCkoL6RW3C"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.nn import functional as F\n",
        "from copy import deepcopy"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SdQmEud3zqSp"
      },
      "source": [
        "This function `measure_unpruned_layer_importances` is designed to calculate importance scores for the attention layers in a model.\n",
        "\n",
        "The basic idea is: if a layer's output is very similar to its input, it might not be doing much important work and could be a candidate for pruning. To check the difference I'm using the cosine similarity."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d4jGSXw-yGYs"
      },
      "outputs": [],
      "source": [
        "def measure_unpruned_layer_importances(pruned_model, tokenizer, input_text):\n",
        "    \"\"\"\n",
        "    Measures and returns importance scores for all unpruned (non-bypassed) layers.\n",
        "    \"\"\"\n",
        "    # PREPARATION\n",
        "    \"\"\"\n",
        "    set the model to evaluation mode to ensure that no gradients\n",
        "    are computed during the forward pass.\n",
        "    \"\"\"\n",
        "    pruned_model.eval()\n",
        "    device = next(pruned_model.parameters()).device\n",
        "\n",
        "    \"\"\"\n",
        "    The provided input text (input_text) is tokenized into tensors\n",
        "    suitable for processing by the model.\n",
        "    \"\"\"\n",
        "    inputs = tokenizer(input_text, return_tensors=\"pt\").to(device)\n",
        "\n",
        "    \"\"\"This will hold tuples of (layer_idx, importance_score)\"\"\"\n",
        "    importance_scores = []\n",
        "\n",
        "    # IDENTIFY UNPRUNED LAYERS & CREATING HOOKS\n",
        "    \"\"\"\n",
        "    We'll register hooks for only layers that are NOT in drop_attn_list\n",
        "    The list of attention layers that have already been pruned,\n",
        "    is stored in a variable in the model's config: pruned_model.config.drop_attn_list.\n",
        "    \"\"\"\n",
        "    unpruned_layer_indices = [\n",
        "        idx for idx in range(len(pruned_model.model.layers))\n",
        "        if idx not in pruned_model.config.drop_attn_list\n",
        "    ]\n",
        "\n",
        "    \"\"\"\n",
        "    Temporary storage for each layer's input/output\n",
        "    We'll store them by layer index\n",
        "    \"\"\"\n",
        "    layer_inputs = {}\n",
        "    layer_outputs = {}\n",
        "\n",
        "    \"\"\"\n",
        "    Create 2 hooks to capture the input and the output of the layers.\n",
        "    These hooks store the inputs and outputs in dictionaries\n",
        "    (layer_inputs and layer_outputs) for later analysis\n",
        "    \"\"\"\n",
        "    #Allows capture the input to the query projection (q_proj)\n",
        "    def q_proj_input_hook(layer_idx):\n",
        "        def _hook(module, module_input):\n",
        "            # module_input can be a tuple depending on PyTorch version\n",
        "            inp = module_input[0] if isinstance(module_input, tuple) else module_input\n",
        "            layer_inputs[layer_idx] = inp.detach().clone()\n",
        "        return _hook\n",
        "\n",
        "    # Allows capture the output from the output projection (o_proj)\n",
        "    def o_proj_output_hook(layer_idx):\n",
        "        def _hook(module, module_input, module_output):\n",
        "            out = module_output[0] if isinstance(module_output, tuple) else module_output\n",
        "            layer_outputs[layer_idx] = out.detach().clone()\n",
        "        return _hook\n",
        "\n",
        "    # Register hooks for each unpruned layer\n",
        "    handles = []\n",
        "    for idx in unpruned_layer_indices:\n",
        "        layer = pruned_model.model.layers[idx]\n",
        "        handles.append(layer.self_attn.q_proj.register_forward_pre_hook(q_proj_input_hook(idx)))\n",
        "        handles.append(layer.self_attn.o_proj.register_forward_hook(o_proj_output_hook(idx)))\n",
        "\n",
        "    # FORWARD PASS\n",
        "    \"\"\"\n",
        "    Single forward pass (no gradient needed)\n",
        "    A single forward pass is performed on the input text.\n",
        "    During this pass, the hooks capture the inputs and outputs of the unpruned layers.\n",
        "    This step is done with torch.no_grad(),\n",
        "    ensuring no gradients are calculated, which saves memory and computation.\n",
        "    \"\"\"\n",
        "    with torch.no_grad():\n",
        "        _ = pruned_model(**inputs)\n",
        "\n",
        "    \"\"\"\n",
        "    The hooks are removed after the forward pass\n",
        "    to avoid memory leaks or interference with subsequent operations.\n",
        "    \"\"\"\n",
        "    for h in handles:\n",
        "        h.remove()\n",
        "\n",
        "\n",
        "    #COMPUTE IMPORTANCE SCORES\n",
        "    \"\"\"\n",
        "    For each unpruned layer, the inputs and outputs are flattened into vectors for comparison.\n",
        "\n",
        "    Cosine Similarity: The similarity between the input and output vectors is\n",
        "    computed using cosine similarity. Layers with outputs that are very similar\n",
        "    to their inputs likely contribute less to the model’s overall computation.\n",
        "\n",
        "    Importance Score: The importance score for each layer is calculated as 1−similarity\n",
        "    A higher score indicates that the layer transforms its input significantly\n",
        "    and is therefore more important to the model's function.\n",
        "    \"\"\"\n",
        "    for idx in unpruned_layer_indices:\n",
        "        if idx in layer_inputs and idx in layer_outputs:\n",
        "            inp = layer_inputs[idx]\n",
        "            out = layer_outputs[idx]\n",
        "\n",
        "            inp_flat = inp.view(inp.size(0), -1)\n",
        "            out_flat = out.view(out.size(0), -1)\n",
        "\n",
        "            similarity = F.cosine_similarity(inp_flat, out_flat, dim=1).mean().item()\n",
        "            importance_score = 1 - similarity\n",
        "            importance_scores.append((idx, importance_score))\n",
        "\n",
        "            print(f\"Layer {idx} importance score: {importance_score:.4f}\")\n",
        "\n",
        "    \"\"\"A list of tuples is returned, where each tuple contains the layer index\n",
        "    and its calculated importance score.\"\"\"\n",
        "    return importance_scores\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V8M8su2qcjg2"
      },
      "source": [
        "The function `bypass_single_layer` is used to disable the attention mechanism of a specific layer in the model without permanently removing or modifying the layer.\n",
        "\n",
        "This is achieved by dynamically overriding the layer’s forward method to bypass its attention computation.\n",
        "\n",
        "As the attention layers are grouped with the MLP Layers we can just remove an attention layer without removing the associated MLP layer. But we can bypass the layer.\n",
        "\n",
        "The bypassed layer skips computationally expensive attention operations, reducing inference time and memory usage.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yblvEOSgyKd5"
      },
      "outputs": [],
      "source": [
        "def bypass_single_layer(pruned_model, layer_idx):\n",
        "    \"\"\"\n",
        "    Modifies the specified layer's forward method so that attention is bypassed.\n",
        "    \"\"\"\n",
        "    layer = pruned_model.model.layers[layer_idx]\n",
        "\n",
        "    # get the list once, while we still have access to the full config\n",
        "    skip = pruned_model.config.drop_attn_list\n",
        "    layer.self_attn.layer_idx = layer_idx\n",
        "\n",
        "    # Store the original forward.\n",
        "    if not hasattr(layer.self_attn, '_original_forward'):\n",
        "        layer.self_attn._original_forward = layer.self_attn.forward\n",
        "\n",
        "    # Set a simple bypass flag directly on the attention layer\n",
        "    layer.self_attn._should_bypass = True\n",
        "\n",
        "    # A new forward that checks the bypass flag\n",
        "    def new_attention_forward(attn, hidden_states, attention_mask=None, position_ids=None,\n",
        "                    past_key_value=None, output_attentions=False, use_cache=False,\n",
        "                    **kwargs):\n",
        "        if attn.layer_idx in skip:\n",
        "            # short-circuit the pruned layer\n",
        "            return (hidden_states, None) if use_cache else (hidden_states, None)\n",
        "        return attn._orig_forward(hidden_states, attention_mask, position_ids,\n",
        "                                  past_key_value, output_attentions, use_cache,\n",
        "                                  **kwargs)\n",
        "    #set new forward method\n",
        "    layer.self_attn.forward = new_attention_forward.__get__(layer.self_attn,\n",
        "                                                  type(layer.self_attn))\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def one_shot_pruning_inplace(model, tokenizer, input_text, num_layers_to_prune):\n",
        "    \"\"\"\n",
        "    Performs pruning on the original model without creating a copy.\n",
        "    \"\"\"\n",
        "    # Save original device (should be CPU in your case)\n",
        "    device = next(model.parameters()).device\n",
        "\n",
        "    # Set up pruning list\n",
        "    if not hasattr(model.config, 'drop_attn_list'):\n",
        "        model.config.drop_attn_list = []\n",
        "\n",
        "    # Measure importance\n",
        "    scores = measure_unpruned_layer_importances(model, tokenizer, input_text)\n",
        "\n",
        "    if len(scores) < num_layers_to_prune:\n",
        "        raise ValueError(\"Requested more layers to prune than exist\")\n",
        "\n",
        "    # Sort and select layers to bypass\n",
        "    scores.sort(key=lambda x: x[1])  # ascending\n",
        "    layers_to_bypass = [idx for idx, _ in scores[:num_layers_to_prune]]  # Fixed syntax error\n",
        "\n",
        "    # Bypass selected layers\n",
        "    for idx in layers_to_bypass:\n",
        "        model.config.drop_attn_list.append(idx)\n",
        "        bypass_single_layer(model, idx)\n",
        "        print(f\"Bypassing layer {idx} with importance score {dict(scores)[idx]:.4f}\")\n",
        "\n",
        "    print(f\"Bypassed layers: {sorted(model.config.drop_attn_list)}\")\n",
        "\n",
        "    return model"
      ],
      "metadata": {
        "id": "5xf5Uq6YJVm_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tOuKyH0MjyJM"
      },
      "source": [
        "## Execute Pruning."
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Disclaimer**\n",
        "\n",
        "I'm using a single illustrative prompt so that the code path is easy to follow. In any research or production setting you must feed hundreds or thousands of diverse prompts before deciding which layers to deactivate"
      ],
      "metadata": {
        "id": "wyI5XiqYyBnc"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HvhSYyE2Rk1H",
        "outputId": "efd0cc2d-b9a2-474e-c415-60a7b487e731"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Layer 0 importance score: 1.0711\n",
            "Layer 1 importance score: 0.9268\n",
            "Layer 2 importance score: 0.9592\n",
            "Layer 3 importance score: 0.9586\n",
            "Layer 4 importance score: 0.9556\n",
            "Layer 5 importance score: 1.0017\n",
            "Layer 6 importance score: 1.0273\n",
            "Layer 7 importance score: 1.0322\n",
            "Layer 8 importance score: 1.1082\n",
            "Layer 9 importance score: 1.1019\n",
            "Layer 10 importance score: 1.0403\n",
            "Layer 11 importance score: 0.9950\n",
            "Layer 12 importance score: 1.0824\n",
            "Layer 13 importance score: 1.0197\n",
            "Layer 14 importance score: 0.9663\n",
            "Layer 15 importance score: 0.9675\n",
            "Layer 16 importance score: 0.9030\n",
            "Layer 17 importance score: 0.9355\n",
            "Layer 18 importance score: 1.0156\n",
            "Layer 19 importance score: 0.7300\n",
            "Layer 20 importance score: 0.8689\n",
            "Layer 21 importance score: 0.9351\n",
            "Layer 22 importance score: 0.8706\n",
            "Layer 23 importance score: 0.7950\n",
            "Layer 24 importance score: 0.8481\n",
            "Layer 25 importance score: 0.8776\n",
            "Layer 26 importance score: 0.7842\n",
            "Layer 27 importance score: 0.9855\n",
            "Bypassing layer 19 with importance score 0.7300\n",
            "Bypassing layer 26 with importance score 0.7842\n",
            "Bypassing layer 23 with importance score 0.7950\n",
            "Bypassing layer 24 with importance score 0.8481\n",
            "Bypassed layers: [19, 23, 24, 26]\n"
          ]
        }
      ],
      "source": [
        "pruned_model = one_shot_pruning_inplace(\n",
        "      model,\n",
        "      tokenizer,\n",
        "       \"Hi I'm a sample text, used to calculate the cosine difference between input and output.\",\n",
        "      num_layers_to_prune=4\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "46_tIN1brKZW"
      },
      "source": [
        "# Test Pruned Models\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kFdCesVcnnrM"
      },
      "source": [
        "Now, let's test the pruned model, which is a Llama-3.2-3B model where I have marked 4 Attention layers to be bypassed."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7oCJ8EYLhJg-",
        "outputId": "e082ce54-3b82-4a53-b802-1eaca0b7c597"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Setting `pad_token_id` to `eos_token_id`:128001 for open-end generation.\n",
            "Setting `pad_token_id` to `eos_token_id`:128001 for open-end generation.\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Run 1:\n",
            "Tokenization time: 0.44 ms\n",
            "Generation time: 1508.45 ms\n",
            "Decoding time: 0.23 ms\n",
            "Total time: 1509.20 ms\n",
            "\n",
            "Run 2:\n",
            "Tokenization time: 0.49 ms\n",
            "Generation time: 1506.98 ms\n",
            "Decoding time: 0.19 ms\n",
            "Total time: 1507.74 ms\n",
            "\n",
            "Average time over 2 runs: 1508.39 ms\n",
            "Generated text: ['Paris is the capital of France and also its largest city. It is also one of the most visited tourist destination worldwide with millions of tourists visiting every year. There are many things to do in Paris including sightseeing tours, shopping malls, museums etc', 'Paris is the capital of France and also its largest city. It is also one of the most visited tourist destination worldwide with millions of tourists visiting every year. There are many things to do in Paris including sightseeing tours, shopping malls, museums etc']\n"
          ]
        }
      ],
      "source": [
        "# Test the pruned model\n",
        "pruned_model = pruned_model.to(device) #Move the model to GPU again.\n",
        "generated = get_output(prompt, pruned_model, num_runs=2)\n",
        "print(f\"Generated text: {generated}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-Ic4ux7goGu_"
      },
      "source": [
        "\n",
        "The execution of this second model is slightly faster than that of the base model, and the generated text is fairly accurate, although some repetition can be noticed towards the end of the sentence."
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Store the Model.\n"
      ],
      "metadata": {
        "id": "H7R_VaIPTQhN"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "new_model_name = 'attnprun-llama-3.2-3B'\n",
        "output_dir = './'+new_model_name\n",
        "if not os.path.exists(output_dir):\n",
        "    os.makedirs(output_dir)\n",
        "\n",
        "pruned_model.save_pretrained(output_dir)\n",
        "tokenizer.save_pretrained(output_dir)\n",
        "#new_config.save_pretrained(output_dir)\n",
        "print(f\"Pruned model saved to {output_dir}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ERXYQsy5TUtn",
        "outputId": "23347781-ef0d-41d2-fe76-abf62b5744f3"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Pruned model saved to ./attnprun-llama-3.2-3B\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# 2. Check that config contains layers to skip\n",
        "from transformers import AutoConfig\n",
        "config = AutoConfig.from_pretrained(output_dir)\n",
        "\n",
        "if hasattr(config, \"drop_attn_list\"):\n",
        "    print(f\"drop_attn_list stored: {config.drop_attn_list}\")\n",
        "else:\n",
        "    print(\"drop_attn_list isn't present.\")\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mdE7o0PETggr",
        "outputId": "8106fbb6-d795-4948-f805-6dd39b511237"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "drop_attn_list stored: [19, 26, 23, 24]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Upload to Hugging Face."
      ],
      "metadata": {
        "id": "hG63t8jqVdOx"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "El proceso de subida de este modelo a Hugging es ligeramente más complejo por que se debe almacenar no tan solo el modelo en si, sino tambien el código de la función _bypass_single_layer. Que como recordarás es la función que se encarga de decidir cuando ejecutar o simplemente bypasear una capa de atención.  "
      ],
      "metadata": {
        "id": "YRDoj5bYKWVh"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from huggingface_hub import HfApi, upload_folder, whoami, create_repo"
      ],
      "metadata": {
        "id": "0bi3zX7FVjwp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Step 1: Get your HF username from the current token\n",
        "username = whoami()[\"name\"]  # Returns a dict like {'name': 'your_username', 'email': ...}\n",
        "username"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "id": "_AnSaJPuWffg",
        "outputId": "78ec43c7-4a31-4e21-be3a-52d6942b6b10"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'oopere'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 17
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Step 2: Define repo name\n",
        "repo_id = f\"{username}/{new_model_name}\""
      ],
      "metadata": {
        "id": "KDuY1p8zWjax"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Step 3: Define path to your model\n",
        "output_dir = \"./\"+new_model_name\n"
      ],
      "metadata": {
        "id": "-MNgDvImWmFa"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The function must be saved in a .py file, but since this notebook runs on Colab, I’ve decided the best approach is to create a cell that generates the file to be uploaded.\n",
        "\n",
        "The file contains the custom class PrunedLlamaForCausalLM, which extends Hugging Face’s LlamaForCausalLM.\n",
        "\n",
        "This custom class calls the base constructor, ensuring that the model's configuration file includes the drop_attn_list, which specifies the layers that should be skipped.\n",
        "\n",
        "The forward function is modified only for the layers that need to be skipped; the rest continue executing their standard forward function.\n"
      ],
      "metadata": {
        "id": "dJSb5z2nLRS-"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "custom_model_code = '''\n",
        "from transformers.models.llama.modeling_llama import LlamaForCausalLM\n",
        "\n",
        "class PrunedLlamaForCausalLM(LlamaForCausalLM):\n",
        "    def __init__(self, config):\n",
        "        super().__init__(config)\n",
        "        if not hasattr(config, \"drop_attn_list\"):\n",
        "            config.drop_attn_list = []\n",
        "\n",
        "        for idx in config.drop_attn_list:\n",
        "            self._bypass_single_layer(idx)\n",
        "\n",
        "    def _bypass_single_layer(pruned_model, layer_idx):\n",
        "        \"\"\"\n",
        "        Modifies the specified layer's forward method so that attention is bypassed.\n",
        "        \"\"\"\n",
        "        layer = pruned_model.model.layers[layer_idx]\n",
        "\n",
        "        # get the list once, while we still have access to the full config\n",
        "        skip = pruned_model.config.drop_attn_list\n",
        "        layer.self_attn.layer_idx = layer_idx\n",
        "\n",
        "        # Store the original forward.\n",
        "        if not hasattr(layer.self_attn, '_original_forward'):\n",
        "            layer.self_attn._original_forward = layer.self_attn.forward\n",
        "\n",
        "        # Set a simple bypass flag directly on the attention layer\n",
        "        layer.self_attn._should_bypass = True\n",
        "\n",
        "        # A new forward that checks the bypass flag\n",
        "        def new_attention_forward(attn, hidden_states, attention_mask=None, position_ids=None,\n",
        "                        past_key_value=None, output_attentions=False, use_cache=False,\n",
        "                        **kwargs):\n",
        "            if attn.layer_idx in skip:\n",
        "                # short-circuit the pruned layer\n",
        "                return (hidden_states, None) if use_cache else (hidden_states, None)\n",
        "            return attn._orig_forward(hidden_states, attention_mask, position_ids,\n",
        "                                      past_key_value, output_attentions, use_cache,\n",
        "                                      **kwargs)\n",
        "        #set new forward method\n",
        "        layer.self_attn.forward = new_attention_forward.__get__(layer.self_attn,\n",
        "                                                      type(layer.self_attn))\n",
        "\n",
        "'''\n",
        "\n",
        "# Define path and write the file\n",
        "os.makedirs(output_dir, exist_ok=True)\n",
        "with open(os.path.join(output_dir, \"modeling_attnprun_llama.py\"), \"w\") as f:\n",
        "    f.write(custom_model_code.strip())\n",
        "\n",
        "print(\"Custom model script modeling_attnprun_llama.py created successfully.\")\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "L_xaG0FDXojc",
        "outputId": "4dc0ee1d-c49f-41f9-ddad-76ff0da044f4"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Custom model script modeling_attnprun_llama.py created successfully.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now the model's configuration file is updated by adding the `auto_map` field, which tells the Transformers library which class to use to construct the model: `modeling_attnprun_llama.PrunedLlamaForCausalLM.`\n"
      ],
      "metadata": {
        "id": "r59Yae3jVMpc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import json\n",
        "import os\n",
        "\n",
        "# Path to the config file\n",
        "config_path = os.path.join(output_dir, \"config.json\")\n",
        "\n",
        "# Load the existing config\n",
        "with open(config_path, \"r\") as f:\n",
        "    config = json.load(f)\n",
        "\n",
        "# Add or update the auto_map section\n",
        "config[\"auto_map\"] = {\n",
        "    \"AutoModelForCausalLM\": \"modeling_attnprun_llama.PrunedLlamaForCausalLM\"\n",
        "}\n",
        "\n",
        "# Optional: ensure the architecture field is aligned\n",
        "config[\"architectures\"] = [\"PrunedLlamaForCausalLM\"]\n",
        "\n",
        "# Save the updated config\n",
        "with open(config_path, \"w\") as f:\n",
        "    json.dump(config, f, indent=2)\n",
        "\n",
        "print(\"config.json updated with auto_map and architecture.\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gBbxbBbmbeRl",
        "outputId": "f5b15a92-2770-47c4-a98e-832f68d2d2af"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "config.json updated with auto_map and architecture.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Time to upload the folder containing the weights, the config file and the new function to HF."
      ],
      "metadata": {
        "id": "0ciXxElsWsQm"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "create_repo(repo_id=repo_id, repo_type=\"model\", exist_ok=True)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 69
        },
        "id": "TTOCOdMhqzo8",
        "outputId": "db1a98d3-bef7-4c93-e376-7f2286435432"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "RepoUrl('https://huggingface.co/oopere/attnprun-llama-3.2-3B', endpoint='https://huggingface.co', repo_type='model', repo_id='oopere/attnprun-llama-3.2-3B')"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 22
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Step 4: Upload the folder to the Hub\n",
        "upload_folder(\n",
        "    folder_path=output_dir,\n",
        "    path_in_repo=\"\",  # Upload everything to root\n",
        "    repo_id=repo_id,\n",
        "    repo_type=\"model\"\n",
        ")\n",
        "\n",
        "print(f\"Model uploaded successfully to https://huggingface.co/{repo_id}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hCUToiTmWp9M",
        "outputId": "51b81151-6533-4bb6-e070-8d43c1e154e5"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "No files have been modified since last commit. Skipping to prevent empty commit.\n",
            "WARNING:huggingface_hub.hf_api:No files have been modified since last commit. Skipping to prevent empty commit.\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Model uploaded successfully to https://huggingface.co/oopere/attnprun-llama-3.2-3B\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Download model from Hugging Face."
      ],
      "metadata": {
        "id": "8Lp3bWMOehLi"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import gc\n",
        "del pruned_model\n",
        "del tokenizer\n",
        "del model\n",
        "\n",
        "# 2. Libera la caché de la GPU\n",
        "torch.cuda.empty_cache()\n",
        "torch.cuda.ipc_collect()  # Opcional, ayuda en Colab\n",
        "\n",
        "# 3. Forza recolección de basura en Python\n",
        "gc.collect()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "g2eDLdnEmapS",
        "outputId": "a2a18d6e-c62d-463a-e3c3-e9d269bf4251"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "573"
            ]
          },
          "metadata": {},
          "execution_count": 24
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!huggingface-cli cache purge --yes"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Iy6GqernoZJx",
        "outputId": "8463581c-85c4-4ddf-a2d9-8e22c6466975"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "usage: huggingface-cli <command> [<args>]\n",
            "huggingface-cli: error: argument {download,upload,repo-files,env,login,whoami,logout,auth,repo,lfs-enable-largefiles,lfs-multipart-upload,scan-cache,delete-cache,tag,version,upload-large-folder}: invalid choice: 'cache' (choose from 'download', 'upload', 'repo-files', 'env', 'login', 'whoami', 'logout', 'auth', 'repo', 'lfs-enable-largefiles', 'lfs-multipart-upload', 'scan-cache', 'delete-cache', 'tag', 'version', 'upload-large-folder')\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The model is downloaded normally from Hugging Face, but you must remember to set `trust_remote_code=True` since the model includes the custom code you previously created and uploaded.\n"
      ],
      "metadata": {
        "id": "UgHQ5-2MW6q2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "model_hf = AutoModelForCausalLM.from_pretrained(\n",
        "    repo_id,\n",
        "    trust_remote_code=True\n",
        ")\n"
      ],
      "metadata": {
        "id": "d4AB09pNemii"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "tokenizer = AutoTokenizer.from_pretrained(repo_id)"
      ],
      "metadata": {
        "id": "oQArVXsSmtp-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model_hf = model_hf.to(device) #Move the model to GPU again.\n",
        "generated = get_output(prompt, model_hf, num_runs=2)\n",
        "print(f\"Generated text: {generated}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2feKDCrPll0E",
        "outputId": "be1d7ecc-817d-4229-b8d0-bb011029e55f"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Setting `pad_token_id` to `eos_token_id`:128001 for open-end generation.\n",
            "Setting `pad_token_id` to `eos_token_id`:128001 for open-end generation.\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Run 1:\n",
            "Tokenization time: 0.74 ms\n",
            "Generation time: 1530.55 ms\n",
            "Decoding time: 0.20 ms\n",
            "Total time: 1531.58 ms\n",
            "\n",
            "Run 2:\n",
            "Tokenization time: 0.53 ms\n",
            "Generation time: 1535.94 ms\n",
            "Decoding time: 0.21 ms\n",
            "Total time: 1536.77 ms\n",
            "\n",
            "Average time over 2 runs: 1534.09 ms\n",
            "Generated text: ['Paris is the capital of France and also its largest city. It is also one of the most visited tourist destination worldwide with millions of tourists visiting every year. There are many things to do in Paris including sightseeing tours, shopping malls, museums etc', 'Paris is the capital of France and also its largest city. It is also one of the most visited tourist destination worldwide with millions of tourists visiting every year. There are many things to do in Paris including sightseeing tours, shopping malls, museums etc']\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5CsaXViPtmft"
      },
      "source": [
        "# Conclusion.\n",
        "Based on the findings in the paper and the results obtained, I believe this type of pruning may work better with larger models where attention layers tend to have redundancy.\n",
        "\n",
        "Since this type of pruning does not alter the model's structure, it does not result in a reduction in its size or the memory required to load it. The main advantage of using this pruning approach is the reduction of computational load during inference, leading to a more efficient model with faster responses and lower resource consumption.\n",
        "\n",
        "Unlike the original paper, which describes \"removing\" selected attention layers but provides limited implementation details, this implementation takes a transparent functional approach by explicitly overriding the `forward` method only in the specified layers. As a result, the model retains its full architecture and parameter set, but selectively skips computations at runtime. This makes the method reversible, modular, and fully compatible with the Hugging Face ecosystem using `trust_remote_code=True`. While both approaches achieve similar computational savings, this one emphasizes clarity, portability, and practical integration.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-XDtwh-dwMHh"
      },
      "source": [
        "# Authors Note.\n",
        "\n",
        "In addition to creating content like this notebook and offering it under the MIT license, I have also contributed to repositories such as those of Hugging Face and Google Gemini.\n",
        "\n",
        "I am especially proud of my book: [Large Language Models: Apply and Implement Strategies for Large Language Models (Apress)(https://amzn.to/3DSepLb).\n",
        "\n",
        "You can find it on both [Amazon](https://amzn.to/3DSepLb) and [Springer](https://link.springer.com/book/10.1007/979-8-8688-0515-8), where they often have good deals on the purchase price.\n",
        "\n",
        "If you take a look and end up purchasing it, keep in mind that you can reach out with any questions via the Discussions section of this same repository or on any of my social media channels. I’ll do my best to respond as quickly as possible."
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "A100",
      "provenance": [],
      "authorship_tag": "ABX9TyM94s3XHmpJteMPW8MWfPBn",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}